import numpy as np
import pandas as pd
from qtorch.strategy import Strategy

class RSIStrategy(Strategy):
    """RSI策略（14日周期，超买70/超卖30）"""
    def __init__(self, window=14, overbought=70, oversold=30):
        self.window = window
        self.overbought = overbought
        self.oversold = oversold
        
    def generate_signals(self, prices):
        delta = prices.diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        avg_gain = gain.rolling(window=self.window).mean()
        avg_loss = loss.rolling(window=self.window).mean()
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
        
        signals = np.where(rsi > self.overbought, -1, 
                         np.where(rsi < self.oversold, 1, 0))
        return signals.astype(int)